

import pickle
import fiftyone.utils.data as foud
import fiftyone.core.metadata as fom
import fiftyone.core.labels as fol
import os

import tqdm

class MMdetLabeledImageDatasetImporter(foud.LabeledImageDatasetImporter):
    """Custom importer for labeled image datasets.

    Args:
        dataset_dir (None): the dataset directory. This may be optional for
            some importers
        shuffle (False): whether to randomly shuffle the order in which the
            samples are imported
        seed (None): a random seed to use when shuffling
        max_samples (None): a maximum number of samples to import. By default,
            all samples are imported
        **kwargs: additional keyword arguments for your importer
    """

    def __init__(
        self,
        dataset_dir=None,
        shuffle=False,
        seed=None,
        max_samples=None,
        threshold=0.02,
        field_name = 'predict',
        **kwargs,
    ):
        super().__init__(
            dataset_dir=dataset_dir,
            shuffle=shuffle,
            seed=seed,
            max_samples=max_samples,
        )
        self.threshold = [threshold,]*10 if isinstance(threshold, float) else threshold
        self.field_name = field_name

    def __len__(self):
        """The total number of samples that will be imported.

        Raises:
            TypeError: if the total number is not known
        """
        # Return the total number of samples in the dataset (if known)
        return len(self.samples)

    def __next__(self):
        """Returns information about the next sample in the dataset.

        Returns:
            an  ``(image_path, image_metadata, label)`` tuple, where

            -   ``image_path``: the path to the image on disk
            -   ``image_metadata``: an
                :class:`fiftyone.core.metadata.ImageMetadata` instances for the
                image, or ``None`` if :meth:`has_image_metadata` is ``False``
            -   ``label``: an instance of :meth:`label_cls`, or a dictionary
                mapping field names to :class:`fiftyone.core.labels.Label`
                instances, or ``None`` if the sample is unlabeled

        Raises:
            StopIteration: if there are no more samples to import
        """
        # Implement loading the next sample in your dataset here
        image_path, image_metadata, detections = next(self.image_iter)
        # print(image_metadata)
        return image_path, image_metadata, {self.field_name: detections}

    @property
    def has_dataset_info(self):
        """Whether this importer produces a dataset info dictionary."""
        return True

    @property
    def has_image_metadata(self):
        """Whether this importer produces
        :class:`fiftyone.core.metadata.ImageMetadata` instances for each image.
        """
        return True

    @property
    def label_cls(self):
        """The :class:`fiftyone.core.labels.Label` class(es) returned by this
        importer.

        This can be any of the following:

        -   a :class:`fiftyone.core.labels.Label` class. In this case, the
            importer is guaranteed to return labels of this type
        -   a list or tuple of :class:`fiftyone.core.labels.Label` classes. In
            this case, the importer can produce a single label field of any of
            these types
        -   a dict mapping keys to :class:`fiftyone.core.labels.Label` classes.
            In this case, the importer will return label dictionaries with keys
            and value-types specified by this dictionary. Not all keys need be
            present in the imported labels
        -   ``None``. In this case, the importer makes no guarantees about the
            labels that it may return
        """
        # Return the appropriate value here
        return {self.field_name: fol.Detections}
    
    def setup(self):
        """Performs any necessary setup before importing the first sample in
        the dataset.

        This method is called when the importer's context manager interface is
        entered, :func:`DatasetImporter.__enter__`.
        """
        pkl_file = self.dataset_dir
        if not os.path.isfile(pkl_file):
            raise Exception("Invalid mmdet prediction pklfile")
        self.samples = []
        with open(pkl_file, 'rb') as f:
            datas = pickle.load(f)
            print(f'found {len(datas)} samples')
            cls_dist = {str(i):0 for i in range(10)}
            for data in tqdm.tqdm(datas):
                img_path = data['img_path']
                # tag = img_path.split('/')[-3]
                image_metadata = fom.ImageMetadata.build_for(img_path)
                ori_shape = data['ori_shape']
                img_shape = data['img_shape']
                detections = []
                for score, label, bbox in zip(data['pred_instances']['scores'], data['pred_instances']['labels'], data['pred_instances']['bboxes']):
                    if score < self.threshold[label.item()]:
                        continue
                    x1, y1, x2, y2 = bbox
                    x1, y1, x2, y2 = x1/ori_shape[1], y1/ori_shape[0], x2/ori_shape[1], y2/ori_shape[0]
                    detections.append(fol.Detection(label=str(label.item()), bounding_box=[x1, y1, x2-x1, y2-y1], confidence=score))
                    cls_dist[str(label.item())] += 1
                fo_detections = fol.Detections(detections=detections,)
                self.samples.append((img_path, image_metadata, fo_detections))
            print(f'loaded {len(self.samples)} samples')
            print(f'class distribution:')
            for k,v in cls_dist.items():
                print(f'{k}: {v}')
            
            
        self.image_iter = iter(self.samples)
        
        self.info = {"dataset_dir": self.dataset_dir, }

    def get_dataset_info(self):
        """Returns the dataset info for the dataset.

        By convention, this method should be called after all samples in the
        dataset have been imported.

        Returns:
            a dict of dataset info
        """
        return self.info

import fiftyone.types as fot

class MMdetImageDataset(fot.LabeledImageDataset):
    def get_dataset_importer_cls(self):
        return MMdetLabeledImageDatasetImporter
    def get_dataset_exporter_cls(self):
        pass